{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Getting Started with the NeMo Agent Toolkit\n",
    "\n",
    "In this notebook, we walk through the basics of using the toolkit, from installation all the way to creating and running your very own custom workflow."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Environment Setup\n",
    "\n",
    "Ensure you meet the following prerequisites:\n",
    "1. Git\n",
    "2. [uv](https://docs.astral.sh/uv/getting-started/installation/)\n",
    "3. NeMo-Agent-Toolkit installed from source following [these instructions](https://github.com/cdgamarose-nv/NeMo-Agent-Toolkit/tree/develop?tab=readme-ov-file#install-from-source)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set API keys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import getpass\n",
    "import os\n",
    "\n",
    "if \"NVIDIA_API_KEY\" not in os.environ:\n",
    "    nvidia_api_key = getpass.getpass(\"Enter your NVIDIA API key: \")\n",
    "    os.environ[\"NVIDIA_API_KEY\"] = nvidia_api_key\n",
    "\n",
    "if \"TAVILY_API_KEY\" not in os.environ:\n",
    "    tavily_api_key = getpass.getpass(\"Enter your Tavily API key: \")\n",
    "    os.environ[\"TAVILY_API_KEY\"] = tavily_api_key\n",
    "\n",
    "if \"OPENAI_API_KEY\" not in os.environ:\n",
    "    openai_api_key = getpass.getpass(\"Enter your OpenAI API key: \")\n",
    "    os.environ[\"OPENAI_API_KEY\"] = openai_api_key"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bringing an Agent into the NeMo-Agent-Toolkit\n",
    "\n",
    "NeMo Agent toolkit works side-by-side and complements any existing agentic framework or memory tool you're using and isn't tied to any specific agentic framework, long-term memory, or data source. This allows you to use your current technology stack - such as LangChain, LlamaIndex, CrewAI, and Microsoft Semantic Kernel, as well as customer enterprise frameworks and simple Python agents - without replatforming.\n",
    "\n",
    "We'll walk you through how to achieve this.\n",
    "\n",
    "To demonstrate this, let's say that you have the following simple langchain agent that answers generic user queries about current events by performing a web search using Tavily. We will show you how to bring this agent into the NeMo-Agent-Toolkit and benefit from the configurability, resuability, and easy user experience.\n",
    "\n",
    "Run the following two cells to create the langchain agent and run it with an example input."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load langchain_sample/langchain_agent.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python langchain_sample/langchain_agent.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Creating a new NeMo-Agent-Toolkit Workflow \n",
    "\n",
    "Bringing this agent into the toolkit requires creating a new workflow and configuring the tools, and so on. A workflow is a self-contained pipeline that orchestrates tools (e.g., custom arithmetic tools, web search, RAG) and one or more LLMs to process user inputs and generate outputs.\n",
    "\n",
    "With our `nat workflow create` sub-command, you can scaffold and register new workflows within seconds. \n",
    "\n",
    "For example, to create an agent called `first_search_agent` in `.tmp/notebooks` you would run the following commands. \n",
    "\n",
    "> Note: The agent in this example has already been created in `examples/notebooks/first_search_agent` directory.\n",
    "\n",
    "```bash\n",
    "mkdir -p $PROJECT_ROOT/.tmp/notebooks\n",
    "nat workflow create --workflow-dir $PROJECT_ROOT/.tmp/notebooks/first_search_agent\n",
    "```\n",
    "\n",
    "Expected Cell Output:\n",
    "```bash\n",
    "Installing workflow 'first_search_agent'...\n",
    "Workflow 'first_search_agent' installed successfully.\n",
    "Workflow 'first_search_agent' created successfully in '/NeMo-Agent-Toolkit/.tmp/notebooks/first_search_agent'.\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The above command:\n",
    "- Creates a new directory similar to `examples/notebooks/first_search_agent`.\n",
    "- Sets up the necessary files and folders.\n",
    "- Installs the new Python package for your workflow.\n",
    "\n",
    "The registration process is built around two main components:\n",
    "1. **A configuration class that inherits from `WorkflowBaseConfig`**\n",
    "    \n",
    "    Configuration classes that inherit from `TypedBaseModel` and `BaseModelRegistryTag` serve as Pydantic-based configuration objects that define both the plugin type identifier and runtime configuration settings for each NeMo Agent toolkit component. Each plugin type (functions, LLMs, embedders, retrievers, memory, front-ends, etc.) has its own base configuration class (e.g., `FunctionBaseConfig`, `LLMBaseConfig`, `EmbedderBaseConfig`) that establishes the plugin category, while concrete implementations specify a unique name parameter that automatically populates the type field for plugin identification. These configuration classes encapsulate runtime parameters as typed Pydantic fields with validation rules, default values, and documentation (e.g., `api_key`, `model_name`, `temperature` for LLM providers, or `uri`, `collection_name`, `top_k` for retrievers), enabling type-safe configuration management, automatic schema generation, and validation across the entire plugin ecosystem.\n",
    "\n",
    "2. **A decorated async function (with `@register_workflow`) that yields a callable response function.**\n",
    "     \n",
    "     A `FunctionInfo` object is a structured representation yielded from functions decorated with `@register_function` that serves as a framework-agnostic wrapper for callable functions in the NeMo Agent Toolkit. This object encapsulates the function's main callable (e.g., `_response_fn`) that will be invoked at runtime, along with its input/output Pydantic schemas for validation, description for documentation, and optional type converters for automatic type transformation. FunctionInfo objects provide a consistent interface that can be dynamically translated into framework-specific representations (e.g., LangChain tools, LlamaIndex functions) at runtime or invoked directly as standard Python async coroutines, enabling seamless integration across different LLM frameworks while maintaining type safety and validation.\n",
    "\n",
    "\n",
    "Once configured, you can run workflows via the command line (`nat run`) or launch them as services (`nat serve`) to handle requests in real time."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Customizing your Workflow"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now its time to define the same langchain agent inside your newly created workflow. This is as simple as making a few code additions to the `first_search_agent_function`.\n",
    "- Add langchain framework wrappers (all this does is indicate which framework you are wrapping your code in which enables profiling the workflow later)\n",
    "- Paste your agent initialization code inside the `first_search_agent_function`\n",
    "- Paste your agent invocation code inside the `_response_fn` function\n",
    "\n",
    "Your final `first_search_agent_function.py` should look like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load first_search_agent/src/nat_first_search_agent/first_search_agent_function.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once you have your workflow registered, you can reference it by its `_type` in a YAML file. \n",
    "\n",
    "For example:\n",
    "\n",
    "```yaml\n",
    "workflow:\n",
    "  _type: first_search_agent\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Running your Workflow\n",
    "\n",
    "The NeMo Agent toolkit provides several ways to run/host an workflow. These are called `front_end` plugins. Some examples are:\n",
    "\n",
    "console: `nat run` (or long version nat start console …). This is useful when performing local testing and debugging. It allows you to pass inputs defined as arguments directly into the workflow. This is show already in the notebook.\n",
    "\n",
    "Fastapi: `nat serve`(or long version nat start fastapi …). This is useful when hosting your workflow as a REST and websockets endpoint.\n",
    "\n",
    "MCP: `nat mcp` (or long version nat start mcp …). This is useful when hosting the workflow and/or any function as an MCP server\n",
    "\n",
    "While these are the built in front-end components, the system is extensible with new user defined front-end plugins.\n",
    "\n",
    "For more info, here is a good resource for using the various plugins from the CLI: [cli.md](https://github.com/NVIDIA/NeMo-Agent-Toolkit/blob/develop/docs/source/reference/cli.md)\n",
    "\n",
    "In order to test your new agent using the console, run:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!nat run --config_file first_search_agent/configs/config.yml --input \"Who is the current Pope?\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As shown above, this will return the same output as your previously created langchain agent."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Runtime Configurations\n",
    "\n",
    "To benefit from the configurability of this toolkit, we can update the configuration object and config file along with the function to use the parameters at runtime.\n",
    "\n",
    "This involves allowing the toolkit to sets up your tools, LLM, and any additional logic like maximum number of historical messages to provide to the agent, maximum number of iterations to run the agent, description of the agent and so on.\n",
    "\n",
    "The toolkit will make use of the `Builder` class to utilize them at runtime."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Your final configuration object should look like this:\n",
    "```python\n",
    "class SecondSearchAgentFunctionConfig(FunctionBaseConfig, name=\"second_search_agent\"):\n",
    "    \"\"\"\n",
    "    NeMo Agent toolkit function template. Please update the description.\n",
    "    \"\"\"\n",
    "    tool_names: list[FunctionRef] = Field(default=[], description=\"List of tool names to use\")\n",
    "    llm_name: LLMRef = Field(description=\"LLM name to use\")\n",
    "    max_history: int = Field(default=10, description=\"Maximum number of historical messages to provide to the agent\")\n",
    "    max_iterations: int = Field(default=15, description=\"Maximum number of iterations to run the agent\")\n",
    "    handle_parsing_errors: bool = Field(default=True, description=\"Whether to handle parsing errors\")\n",
    "    verbose: bool = Field(default=True, description=\"Whether to print verbose output\")\n",
    "    description: str = Field(default=\"\", description=\"Description of the agent\")\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can then replace:\n",
    "```python\n",
    "tool = [search]\n",
    "```\n",
    "with \n",
    "```python\n",
    "tools = builder.get_tools(config.tool_names, wrapper_type=LLMFrameworkEnum.LANGCHAIN)\n",
    "```\n",
    "> **Note**: This allows you to bring in tools from other frameworks like llama index as well and wrap them with langchain since you are implementing your agent in langchain.\n",
    "\n",
    "In a similar way, you can initialize your llm by utilizing the parameters from the configuration object in the following way:\n",
    "```python\n",
    "llm = await builder.get_llm(config.llm_name, wrapper_type=LLMFrameworkEnum.LANGCHAIN)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For each tool or reusable plugin, there are potentially multiple optional parameters with default values that can be overridden. The `nat info components` command can be used to list all available parameters. For example, to list all available parameters for the LLM nim type run:\n",
    "\n",
    "```bash\n",
    "nat info components -t llm_provider -q nim\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Reusing the Inbuilt Tavily Search Function\n",
    "\n",
    "We can also make use of some of many example functions that the toolkit provides for common use cases. In this agent example, rather than reimplementing the tavily search, we will use the inbuilt function for internet search which is built on top of langchain's tavily search API. You can list available functions using the following:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!nat info components -t function -q tavily_internet_search"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This function can be used any number of times in the configuration YAML by specifying the `_type` as `tavily_internet_search`\n",
    "\n",
    "```yaml\n",
    "functions:\n",
    "  my_internet_search:\n",
    "    _type: tavily_internet_search\n",
    "    max_results: 2\n",
    "    api_key: $TAVILY_API_KEY\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Final Code and Configuration\n",
    "The final code for your workflow can be found in [this example](examples/my_agent_workflow/src/nat_my_agent_workflow/my_agent_workflow_function.py)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load first_search_agent/src/nat_first_search_agent/second_search_agent_function.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The final configuration file should resemble the following:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load first_search_agent/configs/config_modified.yml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!nat run --config_file first_search_agent/configs/config_modified.yml --input \"Who is the current Pope?\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### NAT Serve\n",
    "\n",
    "You can also use the `nat serve` sub-command to launch a server and make HTTP requests to the endpoints as shown below. Refer to [this documentation](https://docs.nvidia.com/nemo/agent-toolkit/latest/reference/api-server-endpoints.html) for more information on available endpoints."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash --bg\n",
    "# This will start background nat service and might take a moment to be ready\n",
    "nat serve --config_file first_search_agent/configs/config_modified.yml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "# Issue a request to the background service\n",
    "curl --request POST \\\n",
    "  --url http://localhost:8000/chat \\\n",
    "  --header 'Content-Type: application/json' \\\n",
    "  --data '{\n",
    "    \"messages\": [\n",
    "      {\n",
    "        \"role\": \"user\",\n",
    "        \"content\": \"Who is the current Pope?\"\n",
    "      }\n",
    "    ]\n",
    "}'\n",
    "# Terminate the process after completion\n",
    "pkill -9 nat"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Reusing the Inbuilt ReAct Agent\n",
    "\n",
    "NeMo Agent Toolkit has a reusable react agent function. We can reuse that agent here to simplify the workflow even further."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!nat info components -t function -q react_agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load first_search_agent/configs/config_react_agent.yml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!nat run --config_file first_search_agent/configs/config_react_agent.yml --input \"Who is the current Pope?\""
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
